#include <cstdlib>
#include <cstdio>
#include <cstring>
#include <sstream>

#include "md5.h"
#include "halfproxy_auth.h"

#define BLOCKSIZE	64
#define OUTLEN		16

using namespace std;
using namespace tcetools; 

void check_auth(Socket &s, std::string secret, std::string tag, char myauth_challenge[16]);

void hmac_md5(const char *key, size_t klen, const char *message, size_t mlen, 
		char result[OUTLEN]) {

	char ipad[BLOCKSIZE];
	char opad[BLOCKSIZE];

	for (size_t i = 0; i < BLOCKSIZE; i++) {
		ipad[i] = 0x36;
		opad[i] = 0x5c;
	}

	if (klen > BLOCKSIZE) {
		MD5 h;
		h.update(key, klen);
		h.finalize();
		for (size_t i = 0; i < OUTLEN; i++) {
			ipad[i] ^= h.digest[i];
			opad[i] ^= h.digest[i];
		}
	} else {
		for (size_t i = 0; i < klen; i++) {
			ipad[i] ^= key[i];
			opad[i] ^= key[i];
		}
	}

	MD5 h1;

	h1.update(ipad, BLOCKSIZE);
	h1.update(message, mlen);
	h1.finalize();

	MD5 h2;

	h2.update(opad, BLOCKSIZE);
	h2.update(h1.digest, OUTLEN);
	h2.finalize();

	for (size_t i = 0; i < OUTLEN; i++) result[i] = h2.digest[i];
}

void gen_challenge(char challenge[16]) {
	FILE *fh;

	fh = fopen("/dev/urandom", "r");
	fread(challenge, 1, 16, fh);
	fclose(fh);
}

void gen_response(
			char response[16], char challenge[16], 
			std::string secret) {
	hmac_md5(secret.c_str(), secret.length(), challenge, 16, response);
}

bool flush_http_response(Socket &s) {
	char buf[16];
	
	// Check for correct request line
	s.blocking_unbuffered_read(buf, 15);
	buf[15] = 0;
	if (string(buf).compare("HTTP/1.1 200 OK") != 0) return false;
	
	// Now, we look for /r/n/r/n
	buf[0] = buf[1] = buf[2] = buf[3] = buf[4] = 0;
	while (string(buf).compare("\r\n\r\n") != 0) {
		buf[0] = buf[1];
		buf[1] = buf[2];
		buf[2] = buf[3];
		if (s.raw_read(&buf[3], 1) == 0) return false;
	}
	
	return true;
}

void check_auth_client(IPSocket &s, string secret, string tag) {
	stringstream req;
	
	req << "POST / HTTP/1.1\r\n"
		   "Host: " << s.get_rem_ip() << "\r\n"
		   "Content-Type: application/octet-stream\r\n"
		   "Transfer-Encoding: chunked\r\n"
		   "\r\n";
	
	// Send the HTTP request
	s.blocking_write(req.str().c_str());
	
	// Send the tag
	s.http_chunk_write(tag.c_str(), tag.length());
	
	// Generate a challenge, then send it
	char myauth_challenge[16];
	gen_challenge(myauth_challenge);
	s.http_chunk_write(myauth_challenge, 16);
	
	// Read past HTTP header
	if (!flush_http_response(s)) {
		s.close_sock();
		return;
	}
	
	check_auth(s, secret, tag, myauth_challenge);
}

// Returns true if it passes basic sanity check
bool flush_http_request(IPSocket &s) {
	char buf[16];
	
	// Check for correct request line
	s.blocking_unbuffered_read(buf, 15);
	buf[15] = 0;
	if (string(buf).compare("POST / HTTP/1.1") != 0) return false;
	
	// Now, we look for /r/n/r/n
	buf[0] = buf[1] = buf[2] = buf[3] = buf[4] = 0;
	while (string(buf).compare("\r\n\r\n") != 0) {
		buf[0] = buf[1];
		buf[1] = buf[2];
		buf[2] = buf[3];
		if (s.raw_read(&buf[3], 1) == 0) return false;
	}
	
	return true;
}

void check_auth_server(IPSocket &s, string secret, string tag) {
	const char *resp;
	resp =  "HTTP/1.1 200 OK\r\n"
			"Cache-Control: no-cache\r\n"
			"Content-Type: application/octet-stream\r\n"
			"Transfer-Encoding: chunked\r\n"
			"\r\n";
			
	// Lets clear out the HTTP header
	if (!flush_http_request(s)) {
		s.close_sock();
		return;
	}
	
	// Send the HTTP response
	s.blocking_write(resp, strlen(resp));
	
	// Send the tag
	s.http_chunk_write(tag.c_str(), tag.length());
	
	// Generate a challenge, then send it
	char myauth_challenge[16];
	gen_challenge(myauth_challenge);
	s.http_chunk_write(myauth_challenge, 16);
		
	check_auth(s, secret, tag, myauth_challenge);
}

void check_auth(Socket &s, std::string secret, std::string tag, char myauth_challenge[16]) {
	char myauth_expected[16], myauth_response[16];
	char theirauth_challenge[16], theirauth_response[16];
	
	// Read tag sent to us and verify it's the same
	char *incoming_tag = new char[tag.size() + 1];
	s.http_chunk_read(incoming_tag, tag.size());
	incoming_tag[tag.size()] = 0;
	if (tag.compare(incoming_tag) != 0) {
		s.close_sock();
		return;
	}
	
	// Read their challenge
	s.http_chunk_read(theirauth_challenge, 16);
	
	// Calculate the responses
	gen_response(myauth_expected, myauth_challenge, secret);
	gen_response(theirauth_response, theirauth_challenge, secret);
	
	// Send our response
	s.http_chunk_write(theirauth_response, 16);
	
	// Read their response and verify
	s.http_chunk_read(myauth_response, 16);
	for (int i = 0; i < 16; i++) {
		if (myauth_expected[i] != myauth_response[i]) {
			s.close_sock();
			return;
		}
	}
}


